Summary

Extend the Qwen3-14B device sampling path introduced by #47 to support non-greedy top-k sampling.

For supported requests, a standalone NPU operator reduces full-vocabulary logits to a fixed-width candidate set. The existing host sampler consumes only the candidate values and token IDs, then applies temperature, top-p filtering, distribution validation, and multinomial sampling.

Closes #74.

Paired exact-selector pypto-lib PR: hw-native-sys/pypto-lib#787

What Changed

• Add an executor capability for the maximum device top-k width.
• Add SamplingCandidates to prefill/decode result contracts.
• Automatically enable device candidate selection when temperature > 0 and 0 < top_k <= 32.
• Preserve the existing CPU full-logits fallback for unsupported configurations.
• Compile the Qwen3 selector as a separate L3 callable.
• Reuse one selector for prefill greedy top-1 and non-greedy top-k candidate selection; decode greedy remains in the decode path introduced by Qwen3-14B serving: device-side embedding and sampling #47.
• Allocate reusable prefill/decode candidate buffers and return only active rows.
• Reuse the existing sampler for temperature, top-p, distribution validation, and multinomial selection.
• Wire both offline generation and the serving worker through the same capability-based path.
• Add routing, fallback, candidate sampling, row validation, and operator integration tests.

Scope

• Qwen3-14B on the NPU backend.
• Fixed device candidate capacity of 32; requests with a larger top_k use the CPU fallback.
• The selector remains a standalone operator and is not fused into prefill or decode.
• Final probabilistic sampling remains on the host.

Performance

Final exact top-k: five-run feature comparison

Measured on device 0 with /data/models/Qwen3-14B, prompt 北京故宫是, 128 generated tokens, temperature=0.8, top_k=32, and top_p=1.0. Model initialization and kernel compilation are excluded.

The benchmark isolated the feature before the latest rebase:

• CPU top-k baseline: the measured feature commit's exact parent, pypto-serving@e57e14f.
• Exact NPU top-k: the measured feature commit, pypto-serving@3658ed8.
• Both revisions use pypto-lib@7e7d4cc, the same zm_pypto environment, the same simpler wheel, PTO-ISA checkout, ptoas v0.4.8, model files, and device.
• Baseline and current runs were interleaved on the same device, five runs per revision.

The PR is now rebased as pypto-serving@553252f on upstream/main@1ebd330. The rebase retained the same top-k implementation and only resolved the upstream Qwen BATCH/BATCH_PAD compatibility change; the performance table records the measured pre-rebase revisions above.

The individual 128-token E2E runs were:

• CPU top-k baseline: 16.685, 25.575, 8.991, 9.391, and 9.730 s.
• Exact NPU top-k: 4.385, 4.406, 4.444, 4.488, and 4.281 s.

The CPU path has substantial host-side variance, so both mean and median are reported. Even the conservative median comparison reduces E2E latency by 54.7% and improves throughput by 2.21x. The NPU path is also stable across the five runs, with a 4.281-4.488 s E2E range.

The gain primarily comes from keeping decode full-vocabulary logits device-resident and returning only 32 candidate values and token IDs per active row. Decode kernel time changes only modestly; the dominant full-logits D2H and host full-vocabulary top-k work outside the executor APIs is removed.

Exactness Design

The follow-up exact implementation replaces the original per-512-chunk Top-4 approximation with a provably exact hierarchical reduction:

1. Split 151936 real-vocabulary logits into 74 full 2048-token groups and one 384-token tail group.
2. Compute an exact Top-32 for every group using sort32 and mrgsort.
3. Merge the resulting 75 x 32 = 2400 candidates in a 4096-entry padded buffer and select the exact global Top-32.
4. Compare the NPU output directly against torch.topk(logits[:, :REAL_VOCAB], 32).

The proof is straightforward: if an entry is not in its group's Top-32, at least 32 entries in the same group rank ahead of it, so it cannot belong to the global Top-32. Therefore the union of exact group Top-32 sets contains the exact global Top-32.

The test fixture covers both adversarial distributions:

• All 32 largest values concentrated in one 512-token region, which fails the original Top-4 approximation.
• One of the 32 largest values placed in each of 32 different chunks, which checks token-ID propagation.

The 2048-token grouping also reduces the generated sorting/merge task graph compared with computing Top-32 independently for all 297 512-token chunks. The exact implementation passed the shared-L3 batch-16 warmup and full prefill/decode generation without the 507018 timeout seen with larger task-graph prototypes.

Data Movement Note

For supported decode requests, the decode callable writes full-vocabulary logits to a device scratch tensor, and the standalone top-k selector reads that device tensor directly. DecodeResult.logits is None; only 32 candidate values and 32 token IDs per active row are copied back for the existing host probability sampler.

Prefill still retains its full-logits result contract for fallback. Unsupported configurations, including top_k > 32, continue to use the existing CPU full-logits path.

Validation

• Exact Top-32 NPU output passed direct full-vocabulary torch.topk comparison.
• Greedy selection_k=1 regression passed host argmax comparison.
• Focused rebase validation in zm_pypto: 9 top-k/Qwen routing and compile-interface tests passed.
• tests/test_device_sampling_submission.py: 5 passed after rebase.
• git diff --check: passed in serving and pypto-lib.
• Shared-L3 batch-16 prefill/decode warmup: passed.
• Rebased standalone selector, 100 rounds: 428.0 us mean.
• Five interleaved 128-token exact-parent A/B runs completed successfully on device 0.